package cn.pugle.oj.leetcode;

import cn.pugle.oj.catalog.HeapProblem;
import cn.pugle.oj.catalog.Unknown;

import java.util.ArrayList;
import java.util.Arrays;

/**
 * 最好的方案还是堆, 但是这题我们联系一下快排思想
 *
 * @author tzp
 * @since 2020/10/26
 */
public class LC973 implements HeapProblem {
    class Pair<A, B> {
        A a;
        B b;

        public Pair(A a, B b) {
            this.a = a;
            this.b = b;
        }
    }


    ArrayList<int[]> result = new ArrayList<>();

    public int[][] kClosest(int[][] points, int K) {
        Pair<Double, int[]>[] x = new Pair[points.length];
        for (int i = 0; i < points.length; i++) {
            double sqrt = Math.sqrt(points[i][0] * points[i][0] + points[i][1] * points[i][1]);
            x[i] = new Pair<>(sqrt, points[i]);
        }
        quickSortK(x, 0, x.length - 1, K);
        return result.toArray(new int[0][]);
    }

    /**
     * 从start, end范围, 找前K小的.
     *
     * @param start include
     * @param end   include
     */
    public void quickSortK(Pair<Double, int[]>[] points, int start, int end, int K) {
        if (K == 0) return;
        if (K == 1) {
            double min = Double.MAX_VALUE;
            int[] ele = null;
            for (int i = start; i <= end; i++) {
                if (points[i].a < min) {
                    min = points[i].a;
                    ele = points[i].b;
                }
            }
            result.add(ele);
            return;
        } else {
            double pivot = points[start].a;
            int pivotIndex = start;
            //partition;
            for (int i = start + 1; i <= end; i++) {
                if (points[i].a <= pivot) {
                    pivotIndex++;
                    Pair<Double, int[]> tmp = points[pivotIndex];//swap pivotIndex, i
                    points[pivotIndex] = points[i];
                    points[i] = tmp;
                } else {

                }
            }
            Pair<Double, int[]> tmp = points[pivotIndex];
            points[pivotIndex] = points[start];
            points[start] = tmp;
            int leftNums = pivotIndex - start + 1;
            if (leftNums > K) {//左侧的超过K, 全从左面找
                quickSortK(points, start, pivotIndex, K);
            } else {//左侧的不够, 把左边都加result, 剩下的从右面找
                for (int i = start; i <= pivotIndex; i++) {
                    result.add(points[i].b);
                }
                quickSortK(points, pivotIndex + 1, end, K - leftNums);
            }
        }
    }

    public static void main(String[] args) {
        System.out.println(Arrays.deepToString(new LC973().kClosest(new int[][]{
                {3, 3}, {5, -1}, {-2, 4},
        }, 3)));
    }
}
